#include "../include/Sim3Solver.hpp"
namespace ORBSLAM
{
    Sim3Solver::Sim3Solver(/* args */)
    {
    }

    Sim3Solver::~Sim3Solver()
    {
    }

    Sim3Solver::Sim3Solver(KeyFrame *p_keyframe1, KeyFrame *p_keyframe2, const vector<MapPoint *> &vp_matched12, const bool b_fix_scale, vector<KeyFrame *> vp_matched_mpt_keyframe):mn_iteration(0),mn_best_inliner(0),mb_fix_scale(b_fix_scale)
    {
        // mn_iteration = 0;
        // mn_best_inliner = 0;
        // mb_fix_scale = b_fix_scale;

        bool b_different_keyframe = false;
        if (vp_matched_mpt_keyframe.empty())
        {
            b_different_keyframe = true;
            vp_matched_mpt_keyframe = vector<KeyFrame *>(vp_matched12.size(), p_keyframe2);
        }

        mp_keyframe1 = p_keyframe1;
        mp_keyframe2 = p_keyframe2;

        vector<MapPoint *> vp_keyframe1_mpt = p_keyframe1->Get_vect_mappoints();
        mn_match_num = vp_matched12.size();

        mvp_mpt1.reserve(mn_match_num);
        mvp_mpt2.reserve(mn_match_num);
        mvp_matched12 = vp_matched12;
        mvn_indices1.reserve(mn_match_num);
        mvmat_x3dc1.reserve(mn_match_num);
        mvmat_x3dc2.reserve(mn_match_num);

        const Mat Rcw1 = p_keyframe1->Get_Rcw();
        const Mat tcw1 = p_keyframe1->Get_tcw();
        const Mat Rcw2 = p_keyframe2->Get_Rcw();
        const Mat tcw2 = p_keyframe2->Get_tcw();

        mvn_all_indices.reserve(mn_match_num);

        int index = 0;

        KeyFrame *p_keyframe_m = p_keyframe2;
        for (int i1 = 0; i1 < mn_match_num; i1++)
        {
            if (vp_matched12[i1])
            {
                MapPoint *p_mpt1 = vp_keyframe1_mpt[i1];
                MapPoint *p_mpt2 = vp_matched12[i1];

                if (!p_mpt1)
                {
                    continue;
                }
                if (p_mpt1->Is_bad() || p_mpt2->Is_bad())
                {
                    continue;
                }

                if (b_different_keyframe)
                {
                    p_keyframe_m = vp_matched_mpt_keyframe[i1];
                }
                int index_keyframe1 = p_mpt1->Get_index_in_keyframe(p_keyframe1);
                int index_keyframe2 = p_mpt2->Get_index_in_keyframe(p_keyframe2);

                if (index_keyframe1 < 0 || index_keyframe2 < 0)
                {
                    continue;
                }

                const KeyPoint &kp1 = p_keyframe1->mv_orb_unkeypoints[index_keyframe1];
                const KeyPoint &kp2 = p_keyframe_m->mv_orb_unkeypoints[index_keyframe2];

                const float sigma_square1 = p_keyframe1->mvd_level_sigma2[kp1.octave];
                const float sigma_square2 = p_keyframe2->mvd_level_sigma2[kp2.octave];

                mvf_maxerror1.push_back(9.210 * sigma_square1);
                mvf_maxerror2.push_back(9.210 * sigma_square2);

                mvp_mpt1.push_back(p_mpt1);
                mvp_mpt2.push_back(p_mpt2);
                mvn_indices1.push_back(i1);

                Mat x3d1w = p_mpt1->Get_world_pose();
                mvmat_x3dc1.push_back(Rcw1*x3d1w+tcw1);

                Mat x3d2w = p_mpt2->Get_world_pose();
                mvmat_x3dc2.push_back(Rcw2*x3d2w+tcw2)  ;

                mvn_all_indices.push_back(index);
                index++;
            }
        }

        From_camera_to_image(mvmat_x3dc1,mmat_p1im1);
        From_camera_to_image(mvmat_x3dc2,mmat_p2im2);

        Set_parameters();
    }

    void Sim3Solver::Set_parameters(double probability, int n_inliner, int n_max_iteration)
    {
        md_ransac_prob = probability;
        mn_ransac_min_inliner = n_inliner;
        mn_ransac_max_its = n_max_iteration;

        mn_mpt_num = mvp_mpt1.size();
        mvb_inlineri.resize(mn_mpt_num);
        float epsilon = (float)mn_ransac_min_inliner / mn_mpt_num;

        int n_itetation;

        if (mn_ransac_min_inliner == mn_mpt_num)
        {
            n_itetation = 1;
        }
        else
        {
            n_itetation = ceil(log(1 - md_ransac_prob) / log(1 - pow(epsilon, 3)));
        }

        mn_ransac_max_its = max(1, min(n_itetation, mn_ransac_max_its));
        mn_iteration = 0;
    }

    Mat Sim3Solver::Iterate(int n_ietration, bool &b_nomore, vector<bool> &vb_inliner, int &n_inliner, bool &b_coverge)
    {
        b_nomore = false;
        b_coverge = false;
        vb_inliner = vector<bool>(mn_match_num, false);
        n_inliner = 0;

        if (mn_mpt_num < mn_ransac_min_inliner)
        {
            b_nomore = true;
            return Mat();
        }

        vector<int> vn_available_indices;
        Mat P3dc1i(3, 3, CV_32F);
        Mat P3dc2i(3, 3, CV_32F);

        int n_current_iterations = 0;
        Mat best_sim3;

        while (mn_iteration < mn_ransac_max_its && n_current_iterations < n_ietration)
        {
            mn_iteration++;
            n_current_iterations++;

            vn_available_indices = mvn_all_indices;

            srand((unsigned)time(NULL));
            for (int i = 0; i < 3; i++)
            {
                int randi = rand() % (vn_available_indices.size());
                int idx = vn_available_indices[randi];
                mvmat_x3dc1[idx].copyTo(P3dc1i.col(i));
                mvmat_x3dc2[idx].copyTo(P3dc2i.col(i));
                vn_available_indices[randi] = vn_available_indices.back();
                vn_available_indices.pop_back();
            }

            Compute_sim3(P3dc1i, P3dc2i);
            Check_inline();
            if (mn_inlineri > mn_best_inliner)
            {
                mvb_best_inliner = mvb_inlineri;
                mn_best_inliner = mn_inlineri;
                mmat_best_T12 = mmat_T12i.clone();
                mmat_best_rotation = mmat_r12i.clone();
                mmat_best_translation = mmat_t12i.clone();
                mf_best_scale = mf_s12i;
                if (mn_inlineri > mn_ransac_min_inliner)
                {
                    n_inliner = mn_inlineri;
                    for (int i = 0; i < mn_mpt_num; i++)
                    {
                        if (mvb_inlineri[i])
                        {
                            vb_inliner[mvn_indices1[i]] = true;
                        }
                    }
                    b_coverge = true;
                    return mmat_best_T12;
                }
                else
                {
                    best_sim3 = mmat_best_T12;
                }
            }
        }
        if (mn_iteration >= mn_ransac_max_its)
        {
            b_nomore = true;
        }
        return best_sim3;
    }

    Mat Sim3Solver::Get_estimated_rotation()
    {
        return mmat_best_rotation.clone();
    }

    Mat Sim3Solver::Get_estimated_translation()
    {
        return mmat_best_translation.clone();
    }

    float Sim3Solver::Get_estimated_scale()
    {
        return mf_best_scale;
    }

    void Sim3Solver::Compute_centroid(Mat &P, Mat &Pr, Mat &C)
    {
        reduce(P, C,1, REDUCE_SUM);
        C = C / P.cols;
        for (int i = 0; i < P.cols; i++)
        {
            Pr.col(i) = P.col(i)-C;
        }
    }

    void Sim3Solver::Compute_sim3(Mat &P1, Mat &P2)
    {
        Mat Pr1(P1.size(), P1.type());
        Mat Pr2(P2.size(), P2.type());
        Mat O1(3, 1, Pr1.type());
        Mat O2(3, 1, Pr2.type());

        Compute_centroid(P1, Pr1, O1);
        Compute_centroid(P2, Pr2, O2);

        Mat M = Pr2 * Pr1.t();
        double N11, N12, N13, N14, N22, N23, N24, N33, N34, N44;

        Mat N(4, 4, P1.type());
        float Sxx = M.at<float>(0, 0);
        float Sxy = M.at<float>(0, 1);
        float Sxz = M.at<float>(0, 2);
        float Syx = M.at<float>(1, 0);
        float Syy = M.at<float>(1, 1);
        float Syz = M.at<float>(1, 2);
        float Szx = M.at<float>(2, 0);
        float Szy = M.at<float>(2, 1);
        float Szz = M.at<float>(2, 2);

        N11 = Sxx + Syy + Szz;
        N12 = Syz - Szy;
        N13 = Szx - Sxz;
        N14 = Sxy - Syx;
        N22 = Sxx - Syy - Szz;
        N23 = Sxy + Syx;
        N24 = Szx + Sxz;
        N33 = -Sxx + Syy - Szz;
        N34 = Syz + Szy;
        N44 = -Sxx - Syy + Szz;

        N = (Mat_<float>(4, 4) << N11, N12, N13, N14, N12, N22, N23, N24, N13, N23, N33, N34, N14, N24, N34, N44);
        Mat eval, evec;
        eigen(N, eval, evec);

        Mat vec(1, 3, evec.type());
        (evec.row(0).colRange(1, 4)).copyTo(vec);
        double ang = atan2(norm(vec), evec.at<float>(0, 0));
        vec = 2 * ang * vec / norm(vec);
        mmat_r12i.create(3, 3, P1.type());
        Rodrigues(vec, mmat_r12i);
        Mat P3 = mmat_r12i * Pr2;
        if (!mb_fix_scale)
        {
            // TODO
            double numerator = Pr1.dot(Pr1);
            double denominator = Pr2.dot(Pr2);
            mf_s12i = sqrt(numerator / denominator);
        }
        else
        {
            mf_s12i = 1.0f;
        }

        mmat_t12i.create(1, 3, P1.type());
        mmat_t12i = O1 - mf_s12i * mmat_r12i * O2;
        mmat_T12i = Mat::eye(4, 4, P1.type());
        Mat sR = mf_s12i * mmat_r12i;

        sR.copyTo(mmat_T12i.rowRange(0, 3).colRange(0, 3));
        mmat_t12i.copyTo(mmat_T12i.rowRange(0, 3).col(3));

        mmat_T21i = Mat::eye(4, 4, P1.type());
        Mat sRinv = (1.0 / mf_s12i) * mmat_r12i.t();
        sRinv.copyTo(mmat_T21i.rowRange(0, 3).colRange(0, 3));
        Mat tinv = -sRinv * mmat_t12i;
        tinv.copyTo(mmat_T21i.rowRange(0, 3).col(3));
    }
    
    void Sim3Solver::Check_inline()
    {
        vector<Mat> vp1im2,vp2im1;
        Project(mvmat_x3dc2,vp2im1,mmat_T12i);
        Project(mvmat_x3dc1,vp1im2,mmat_T21i);

        mn_inlineri = 0;

        for(int i=0;i<mmat_p1im1.size();i++){
            Mat dist1 = mmat_p1im1[i]-vp2im1[i];
            Mat dist2 = mmat_p2im2[i]-vp1im2[i];

            const float err1 = dist1.dot(dist1);
            const float err2 = dist2.dot(dist2);

            if(err1<mvf_maxerror1[i]&&err2<mvf_maxerror2[i]){
                mvb_inlineri[i]  =true;
                mn_inlineri++;
            }else{
                mvb_inlineri[i] = false;
            }
        }
        
    }
    
    void Sim3Solver::Project(const vector<Mat> &vp3dw,vector<Mat> &vp2d,Mat Tcw)
    {
        Mat Rcw = Tcw.rowRange(0,3).colRange(0,3);
        Mat tcw = Tcw.rowRange(0,3).col(3);
        vp2d.clear();
        vp2d.reserve(vp3dw.size());

        for(int i=0;i<vp3dw.size();i++){
            Mat p3dc = Rcw*vp3dw[i]+tcw;
            const float invz = 1/(p3dc.at<float>(2));
            const float x = p3dc.at<float>(0);
            const float y = p3dc.at<float>(1);
            const float z = p3dc.at<float>(2);
            vp2d.push_back(Converter::P3d_to_2dmat(Point3f(x,y,z)));
        }
    }
    
    void Sim3Solver::From_camera_to_image(const vector<Mat> &vp3dc,vector<Mat> &vp2d)
    {
        vp2d.clear();
        vp2d.reserve(vp3dc.size());
        for(int i=0;i<vp3dc.size();i++){
            const float invz = 1/(vp3dc[i].at<float>(2));
            const float x = vp3dc[i].at<float>(0);
            const float y = vp3dc[i].at<float>(1);
            const float z = vp3dc[i].at<float>(2);
            vp2d.push_back(Converter::P3d_to_2dmat(Point3f(x,y,z)));
        }
    }
} // namespace ORBSLAM
